Add partner observations by mpragnay · Pull Request #361 · Emerge-Lab/PufferDrive

mpragnay · 2026-03-22T23:26:42Z

Neural network takes as input only the partner observations(within a radius)
To ensure the nearest partners are always in the observations, we add partners by order of distance to the ego agent. Also added a partner_obs_coverage metric for seeing the percentage of partners actually being seen by the ego agent.
EGO POV renders also only outputs the partner observations

pufferlib/config/ocean/drive.ini

pufferlib/ocean/drive/drive.h

…o align with 3.0_beta

…gent

Copilot

Pull request overview

This PR updates the Ocean Drive environment observation pipeline to include only nearby “partner” agents (within a configurable radius), selecting the nearest partners first, and adds logging/rendering support for this new observation scheme.

Changes:

Added partner_obs_radius configuration/plumbing from INI + Python kwargs into the C env.
Switched partner observation sizing from (MAX_AGENTS - 1) to a fixed MAX_PARTNER_OBSERVATIONS, and implemented nearest-within-radius partner selection plus a partner_obs_coverage metric.
Updated visualization/EGO POV rendering behavior to reflect the new partner-observation radius.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
pufferlib/ocean/env_config.h	Adds `partner_obs_radius` to INI config parsing.
pufferlib/ocean/env_binding.h	Exposes `MAX_PARTNER_OBSERVATIONS` to Python (and drops `MAX_AGENTS`).
pufferlib/ocean/drive/visualize.c	Passes `partner_obs_radius` into visualization env initialization.
pufferlib/ocean/drive/drivenet.h	Updates network partner dimensioning to `MAX_PARTNER_OBSERVATIONS`.
pufferlib/ocean/drive/drive.py	Adds `partner_obs_radius` param/validation and threads it into C init; updates partner obs sizing.
pufferlib/ocean/drive/drive.h	Implements partner selection within radius + nearest ordering; adds `partner_obs_coverage` metric; updates obs sizing.
pufferlib/ocean/drive/drive.c	Minor demo/test updates; threads `partner_obs_radius` into demo env.
pufferlib/ocean/drive/binding.c	Adds `partner_obs_radius` kwarg plumbing into env init + logging.
pufferlib/config/ocean/drive.ini	Adds `partner_obs_radius` to default Drive INI.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

pufferlib/ocean/drive/drive.h

Copilot · 2026-04-05T14:22:41Z

pufferlib/ocean/drive/drive.py

    def _validate_agent_dimensions(self):
        if self.max_agents_per_env < self.min_agents_per_env:
            raise ValueError(
                f"max_agents_per_env ({self.max_agents_per_env}) must be >= min_agents_per_env ({self.min_agents_per_env})"
            )
-        if self.max_agents_per_env > binding.MAX_AGENTS:
-            # TODO: Check needs to be removed once MAX_PARTNER_OBS deprecates MAX_AGENTS
-            raise ValueError(
-                f"max_agents_per_env ({self.max_agents_per_env}) cannot exceed MAX_AGENTS ({binding.MAX_AGENTS}) defined in C code."
-            )
        if self.spawn_width_min < 0.0:


_validate_agent_dimensions() no longer enforces max_agents_per_env <= MAX_AGENTS (the C compile-time limit used for fixed-size buffers like AgentDistance candidates[MAX_AGENTS] and other MAX_AGENTS-sized arrays). With this check removed, configurations can set max_agents_per_env above the C limit, leading to memory corruption/stack overflow when INIT_VARIABLE_AGENT_NUMBER spawns more than MAX_AGENTS agents. Until the C side is fully updated to support >MAX_AGENTS safely, restore an upper-bound validation (or clamp) consistent with the actual C limits.

Copilot · 2026-04-05T14:22:41Z

pufferlib/ocean/env_binding.h

    // Make constants accessible from Python
    PyModule_AddIntConstant(m, "MAX_ROAD_SEGMENT_OBSERVATIONS", MAX_ROAD_SEGMENT_OBSERVATIONS);
-    PyModule_AddIntConstant(m, "MAX_AGENTS", MAX_AGENTS);
+    PyModule_AddIntConstant(m, "MAX_PARTNER_OBSERVATIONS", MAX_PARTNER_OBSERVATIONS);


PyInit_binding() no longer exports the MAX_AGENTS constant to Python (it was replaced with MAX_PARTNER_OBSERVATIONS). Even though current in-repo code no longer references binding.MAX_AGENTS, this is a breaking API change for any external consumers/config validation that relied on it. Consider keeping MAX_AGENTS exported alongside MAX_PARTNER_OBSERVATIONS (or documenting the deprecation and updating downstream callers accordingly).

Suggested change

PyModule_AddIntConstant(m, "MAX_PARTNER_OBSERVATIONS", MAX_PARTNER_OBSERVATIONS);

PyModule_AddIntConstant(m, "MAX_PARTNER_OBSERVATIONS", MAX_PARTNER_OBSERVATIONS);

PyModule_AddIntConstant(m, "MAX_AGENTS", MAX_PARTNER_OBSERVATIONS);

Copilot · 2026-04-05T14:22:41Z

pufferlib/ocean/drive/drive.h

+        obs_idx += 8;
+    }
+
+    // Pad remaining partner obs with zero
+    int remaining_partner_obs = (MAX_PARTNER_OBSERVATIONS - cars_seen) * 8;


compute_partner_observations() hard-codes partner feature width as 8 in several places (obs_idx += 8, (…)* 8). This will silently break if PARTNER_FEATURES changes. Prefer using PARTNER_FEATURES consistently for indexing/padding to keep observation layout self-consistent.

Suggested change

obs_idx += 8;

}

// Pad remaining partner obs with zero

int remaining_partner_obs = (MAX_PARTNER_OBSERVATIONS - cars_seen) * 8;

obs_idx += PARTNER_FEATURES;

}

// Pad remaining partner obs with zero

int remaining_partner_obs = (MAX_PARTNER_OBSERVATIONS - cars_seen) * PARTNER_FEATURES;

eugenevinitsky · 2026-04-05T15:33:31Z

pufferlib/ocean/drive/drive.c

    // Weights* weights = load_weights("resources/drive/puffer_drive_weights.bin");
    Weights *weights = load_weights("puffer_drive_weights.bin");
-    DriveNet *net = init_drivenet(weights, num_agents, CLASSIC, 0);
+    int reward_conditioning = 0;


what're these changes about?

Was just some minor fixes for the demo, I can do it in a separate PR too!

eugenevinitsky · 2026-04-05T15:35:12Z

pufferlib/ocean/drive/drive.h

+        obs[obs_idx] = rel_x * 0.02f;
+        obs[obs_idx + 1] = rel_y * 0.02f;
+        obs[obs_idx + 2] = dz * 0.02f;
+        obs[obs_idx + 3] = partner->sim_width / MAX_VEH_WIDTH;


move to constants file

eugenevinitsky · 2026-04-05T15:35:43Z

pufferlib/ocean/drive/drive.h

+    float cos_heading = cosf(ego_entity->sim_heading);
+    float sin_heading = sinf(ego_entity->sim_heading);


This is a little puzzling. The radius is not dependent upon the frame so having any code depend on which way the ego is facing seems odd

Okay, figured out why they're there but I think code-wise it makes sense for this to be closer to where it's used

eugenevinitsky · 2026-04-05T15:37:58Z

pufferlib/ocean/drive/drive.h

+        float coverage = compute_partner_observations(env, obs, i, obs_idx);
+        env->logs[i].partner_obs_coverage += coverage;
+
+        obs_idx += MAX_PARTNER_OBSERVATIONS * PARTNER_FEATURES;


eugenevinitsky · 2026-04-05T15:40:16Z

pufferlib/ocean/drive/drive.h

+        obs[obs_idx + 5] = other_cos * cos_heading + other_sin * sin_heading;
+        obs[obs_idx + 6] = other_sin * cos_heading - other_cos * sin_heading;


not for now but we really should pull all the transformer code into one place so we can't mess it up anywhere

mpragnay force-pushed the pragnay/add_partner_obs_3.0 branch from 3cd4cdd to becce94 Compare March 23, 2026 01:17

eugenevinitsky reviewed Mar 23, 2026

View reviewed changes

pufferlib/config/ocean/drive.ini Outdated Show resolved Hide resolved

eugenevinitsky reviewed Mar 23, 2026

View reviewed changes

pufferlib/ocean/drive/drive.h Outdated Show resolved Hide resolved

mpragnay added 3 commits March 24, 2026 16:31

Added Partner Obs, partners only render in obs_view, demo bug fixes t…

12231cf

…o align with 3.0_beta

Limit Partner Obs correctly

203456e

Handled C neural network for renders

a16d07b

mpragnay force-pushed the pragnay/add_partner_obs_3.0 branch from 3f1364e to a16d07b Compare March 24, 2026 20:34

mpragnay added 7 commits March 26, 2026 20:08

Merge branch '3.0' into pragnay/add_partner_obs_3.0

6efea82

Changed default partner obs radius

6cb6894

Merge branch '3.0' into pragnay/add_partner_obs_3.0

97ce54f

Merge branch '3.0' into pragnay/add_partner_obs_3.0

4a79515

Removed hard coded constants, error handling

b1f0674

Added Partner Obs coverage metric

957b96a

Sorted Partner Obs to prevent skipping of critical partners for ego a…

7641b23

…gent

mpragnay marked this pull request as ready for review April 5, 2026 14:16

Copilot AI review requested due to automatic review settings April 5, 2026 14:16

Copilot started reviewing on behalf of mpragnay April 5, 2026 14:16 View session

Copilot AI reviewed Apr 5, 2026

View reviewed changes

eugenevinitsky reviewed Apr 5, 2026

View reviewed changes

mpragnay changed the title ~~[WIP]Pragnay/add partner obs 3.0~~ Add partner observations Apr 5, 2026

mpragnay added 2 commits April 5, 2026 18:15

Minor fixes

732eabb

Adding normalization constants

d26be7c

eugenevinitsky approved these changes Apr 6, 2026

View reviewed changes

mpragnay merged commit 6f4eee6 into 3.0 Apr 6, 2026
10 checks passed

	PyModule_AddIntConstant(m, "MAX_PARTNER_OBSERVATIONS", MAX_PARTNER_OBSERVATIONS);
	PyModule_AddIntConstant(m, "MAX_PARTNER_OBSERVATIONS", MAX_PARTNER_OBSERVATIONS);
	PyModule_AddIntConstant(m, "MAX_AGENTS", MAX_PARTNER_OBSERVATIONS);

		float cos_heading = cosf(ego_entity->sim_heading);
		float sin_heading = sinf(ego_entity->sim_heading);

		obs[obs_idx + 5] = other_cos * cos_heading + other_sin * sin_heading;
		obs[obs_idx + 6] = other_sin * cos_heading - other_cos * sin_heading;

Conversation

mpragnay commented Mar 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Copilot AI Apr 5, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 5, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 5, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

mpragnay commented Mar 22, 2026 •

edited

Loading